package dd.lo.classifier;

//import org.datavec.image.loader.NativeImageLoader;
//import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
//import org.deeplearning4j.eval.Evaluation;
//import org.deeplearning4j.nn.api.Model;
//import org.deeplearning4j.nn.api.OptimizationAlgorithm;
//import org.deeplearning4j.nn.conf.LearningRatePolicy;
//import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
//import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
//import org.deeplearning4j.nn.conf.Updater;
//import org.deeplearning4j.nn.conf.inputs.InputType;
//import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
//import org.deeplearning4j.nn.conf.layers.DenseLayer;
//import org.deeplearning4j.nn.conf.layers.OutputLayer;
//import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
//import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
//import org.deeplearning4j.nn.weights.WeightInit;
//import org.deeplearning4j.optimize.api.IterationListener;
//import org.deeplearning4j.util.ModelSerializer;
//import org.nd4j.linalg.activations.Activation;
//import org.nd4j.linalg.api.ndarray.INDArray;
//import org.nd4j.linalg.dataset.DataSet;
//import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
//import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
//import org.nd4j.linalg.factory.Nd4j;
//import org.nd4j.linalg.learning.config.Nesterovs;
//import org.nd4j.linalg.lossfunctions.LossFunctions;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.*;

public class MnistClassifier {
    private static final int N_OUTCOMES = 10;

//    static class ScoreIterationCustomListener implements IterationListener {
//        private int printIterations = 10;
//        private boolean invoked = false;
//        private long iterCount = 0;
//
//        /**
//         * @param printIterations    frequency with which to print scores (i.e., every printIterations parameter updates)
//         */
//        public ScoreIterationCustomListener(int printIterations) {
//            this.printIterations = printIterations;
//        }
//
//        /** Default constructor printing every 10 iterations */
//        public ScoreIterationCustomListener() {}
//
//        @Override
//        public boolean invoked() {
//            return invoked;
//        }
//
//        @Override
//        public void invoke() {
//            this.invoked = true;
//        }
//
//        @Override
//        public void iterationDone(Model model, int iteration) {
//            if (printIterations <= 0)
//                printIterations = 1;
//            if (iterCount % printIterations == 0) {
//                invoke();
//                double result = model.score();
//                System.out.printf("Score at iteration %d is %f\n", iterCount, result);
//            }
//            iterCount++;
//        }
//    }
//
//    private MnistReader trainImgReader;
//    private MnistReader trainLabelReader;
//    private MnistReader testImgReader;
//    private MnistReader testLabelReader;
//
//    private boolean dataSetReady = false;
//
//    public MnistReader getTrainImgReader() {
//        return trainImgReader;
//    }
//
//    public MnistReader getTrainLabelReader() {
//        return trainLabelReader;
//    }
//
//    public MnistReader getTestImgReader() {
//        return testImgReader;
//    }
//
//    public MnistReader getTestLabelReader() {
//        return testLabelReader;
//    }
//
//    private void readDataSet(String folderPath, boolean train, boolean test) throws IOException {
//        if (dataSetReady) return;
//        if (train) {
//            String trainLabelFilePath = folderPath + "/train-labels-idx1-ubyte.gz";
//            File trainLabelFile = new File(trainLabelFilePath);
//            if (!trainLabelFile.exists())
//                throw new FileNotFoundException(String.format("训练集标签文件未能找到（%s）", trainLabelFilePath));
//            trainLabelReader = new MnistReader(trainLabelFile.getAbsolutePath(), MnistReader.TYPE_LABEL);
//            String trainImgFilePath = folderPath + "/train-images-idx3-ubyte.gz";
//            File trainImgFile = new File(trainImgFilePath);
//            if (!trainImgFile.exists())
//                throw new FileNotFoundException(String.format("训练集图片文件未能找到（%s）", trainImgFilePath));
//            trainImgReader = new MnistReader(trainImgFile.getAbsolutePath(), MnistReader.TYPE_IMAGE);
//        }
//        if (test) {
//            String testLabelFilePath = folderPath + "/t10k-labels-idx1-ubyte.gz";
//            File testLabelFile = new File(testLabelFilePath);
//            if (!testLabelFile.exists())
//                throw new FileNotFoundException(String.format("测试集标签文件未能找到（%s）", testLabelFilePath));
//            testLabelReader = new MnistReader(testLabelFile.getAbsolutePath(), MnistReader.TYPE_LABEL);
//            String testImgFilePath = folderPath + "/t10k-images-idx3-ubyte.gz";
//            File testImgFile = new File(testImgFilePath);
//            if (!testImgFile.exists())
//                throw new FileNotFoundException(String.format("测试集图片文件未能找到（%s）", testImgFilePath));
//            testImgReader = new MnistReader(testImgFile.getAbsolutePath(), MnistReader.TYPE_IMAGE);
//        }
//        dataSetReady = true;
//    }
//
//    private DataSetIterator getDataSetIterator(MnistReader imgReader, MnistReader labelReader) throws IOException {
//        int height = imgReader.getHeight(), width = imgReader.getWidth(), n = imgReader.size();
//        ImagePreProcessingScaler scalar = new ImagePreProcessingScaler(0,1);
//        INDArray input = Nd4j.create(new int[]{n, height * width});
//        INDArray output = Nd4j.create(new int[]{n, N_OUTCOMES});
//        for (int i = 0; i < n; ++i) {
//            byte[] labelData = labelReader.get(i);
//            int label = labelData[0] & 0xFF;
//            //Transform image data into image file byte stream
//            byte[] imageData = imgReader.get(i);
//            BufferedImage image = new BufferedImage(imgReader.getWidth(), imgReader.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
//            int idx = 0;
//            for (int y = 0; y < height; ++y) {
//                for (int x = 0; x < width; ++x) {
//                    image.setRGB(x, y, imageData[idx] & 0xFF);
//                    ++idx;
//                }
//            }
//            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
//            ImageIO.write(image, "jpeg", outputStream);
//            BufferedInputStream inputStream = new BufferedInputStream(new ByteArrayInputStream(outputStream.toByteArray()));
//            NativeImageLoader nil = new NativeImageLoader(height, width);
//            INDArray img = nil.asRowVector(inputStream);
//            scalar.transform(img);
//            input.putRow(i, img);
//            output.put(i, label, 1.0);
//        }
//        //Joining input and output matrices into a dataset
//        DataSet dataSet = new DataSet(input, output);
//        //Convert the dataset into a list
//        List<DataSet> listDataSet = dataSet.asList();
//        //Shuffle content of list randomly
//        Collections.shuffle(listDataSet, new Random(System.currentTimeMillis()));
//        int batchSize = 10;
//        //Build and return a dataset iterator
//        return new ListDataSetIterator<>(listDataSet, batchSize);
//    }
//
//    private void buildModel(DataSetIterator trainDsi, DataSetIterator testDsi, String outputFolderPath) {
//        int rngSeed = 1234;
//        int height = trainImgReader.getHeight(), width = trainImgReader.getWidth(), channels = 1;
//        long t0 = System.currentTimeMillis();
//        System.out.println("Build Model...");
//        //神经网络示例1：https://ibjects.medium.com/image-classification-neural-network-tutorial-getting-started-with-deep-learning-for-java-dl4j-b36372aad656
//        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
//                .seed(rngSeed)
//                .updater(new Nesterovs(0.006, 0.9))
//                .l2(0.0001).list()
//                .layer(0, new DenseLayer.Builder().nIn(height * width).nOut(1000).activation(Activation.RELU).weightInit(WeightInit.XAVIER).build())
//                .layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).nIn(1000).nOut(N_OUTCOMES).activation(Activation.SOFTMAX).weightInit(WeightInit.XAVIER).build())
//                .build();
//        //神经网络示例2：http://www.uml.org.cn/ai/201903063.asp
//        //设定动态改变学习速率的策略，key表示小批量迭代到几次
////        Map<Integer, Double> lrSchedule = new HashMap<>();
////        lrSchedule.put(0, 0.06);
////        lrSchedule.put(200, 0.05);
////        lrSchedule.put(600, 0.028);
////        lrSchedule.put(800, 0.0060);
////        lrSchedule.put(1000, 0.001);
////        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
////                .seed(rngSeed)
////                .iterations(1)
////                .regularization(true)
////                .l2(0.0005)
////                .learningRate(0.01)
////                .learningRateDecayPolicy(LearningRatePolicy.Schedule)
////                .learningRateSchedule(lrSchedule)
////                .weightInit(WeightInit.XAVIER)
////                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
////                .updater(Updater.NESTEROVS)
////                .list()
////                .layer(0, new ConvolutionLayer.Builder(5, 5).nIn(channels).stride(1, 1).nOut(20).activation(Activation.IDENTITY).build())
////                .layer(1, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX).kernelSize(2, 2).stride(2, 2).build())
////                .layer(2, new ConvolutionLayer.Builder(5, 5).stride(1, 1).nOut(50).activation(Activation.IDENTITY).build())
////                .layer(3, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX).kernelSize(2, 2).stride(2, 2).build())
////                .layer(4, new DenseLayer.Builder().activation(Activation.RELU).nOut(500).build())
////                .layer(5, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).nOut(N_OUTCOMES).activation(Activation.SOFTMAX).build())
////                .setInputType(InputType.convolutionalFlat(height, width, channels))
////                .backprop(true)
////                .pretrain(false)
////                .build();
//        MultiLayerNetwork model = new MultiLayerNetwork(conf);
//        model.init();
//        //Print score every 500 interaction
//        model.setListeners(new ScoreIterationCustomListener(50));
//        System.out.println("Train Model...");
//        model.fit(trainDsi);
//        //Evaluation
//        System.out.println("Evaluating Model...");
//        Evaluation eval = model.evaluate(testDsi);
//        System.out.println(eval.stats());
//        long t1 = System.currentTimeMillis();
//        double t = (double) (t1 - t0)/1000.0;
//        System.out.printf("\nTotal time: %f seconds\n", t);
//        File outputFile = new File(outputFolderPath + "/minist-model.zip");
//        boolean exist = outputFile.exists();
//        if (!outputFile.exists()) {
//            try {
//                exist = outputFile.createNewFile();
//            } catch (IOException e) {
//                e.printStackTrace();
//            }
//        }
//        if (exist) {
//            try {
//                ModelSerializer.writeModel(model, outputFile, true);
//            } catch (IOException e) {
//                e.printStackTrace();
//            }
//        }
//    }
//
//    public void trainModel(String folderPath, String outputFolderPath) throws IOException {
//        readDataSet(folderPath, true, true);
//        DataSetIterator trainDsi = getDataSetIterator(trainImgReader, trainLabelReader);
//        DataSetIterator testDsi = getDataSetIterator(testImgReader, testLabelReader);
//        buildModel(trainDsi, testDsi, outputFolderPath);
//    }
//
//    public void loadModel(String dataSetFolderPath, String modelPath) throws IOException {
//        MultiLayerNetwork model = ModelSerializer.restoreMultiLayerNetwork(new File(modelPath));
//        readDataSet(dataSetFolderPath, true, true);
//        DataSetIterator testDsi = getDataSetIterator(testImgReader, testLabelReader);
//        Evaluation eval = model.evaluate(testDsi);
//        System.out.println(eval.stats());
//    }
//
//    public void predict(String dataSetFolderPath, String modelPath) throws IOException {
//        MultiLayerNetwork model = ModelSerializer.restoreMultiLayerNetwork(new File(modelPath));
//        readDataSet(dataSetFolderPath, false, true);
//        int height = testImgReader.getHeight(), width = testImgReader.getWidth();
//        int i = (int) (Math.random() * testImgReader.size());
//        byte[] labelData = testLabelReader.get(i);
//        int label = labelData[0] & 0xFF;
//        byte[] imageData = testImgReader.get(i);
//        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);
//        int idx = 0;
//        for (int y = 0; y < height; ++y) {
//            for (int x = 0; x < width; ++x) {
//                image.setRGB(x, y, imageData[idx] & 0xFF);
//                ++idx;
//            }
//        }
//        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
//        ImageIO.write(image, "jpeg", outputStream);
//        BufferedInputStream inputStream = new BufferedInputStream(new ByteArrayInputStream(outputStream.toByteArray()));
//        NativeImageLoader nil = new NativeImageLoader(height, width);
//        INDArray img = nil.asRowVector(inputStream);
//        int[] predictResult = model.predict(img);
//        System.out.printf("test %d, predict %d\n", label, predictResult[0]);
//    }
//
//    public static void main(String[] args) throws IOException {
//        MnistClassifier mnistClassifier = new MnistClassifier();
////        mnistClassifier.trainModel("/Users/kwok/Downloads/MNIST", "/Users/kwok/Downloads/MNIST_Classifier");
////        mnistClassifier.loadModel("/Users/kwok/Downloads/MNIST", "/Users/kwok/Downloads/MNIST_Classifier/minist-model.zip");
//        mnistClassifier.predict("/Users/kwok/Downloads/MNIST", "/Users/kwok/Downloads/MNIST_Classifier/minist-model.zip");
//    }
}
